#!/bin/bash

# DNS配置脚本
# 提供预设DNS选项和自定义DNS配置功能

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# DNS预设配置
declare -A DNS_PRESETS
DNS_PRESETS["谷歌"]="8.8.8.8,8.8.4.4"
DNS_PRESETS["Cloudflare"]="1.1.1.1,1.0.0.1"
DNS_PRESETS["腾讯云"]="119.29.29.29,182.254.116.116"
DNS_PRESETS["阿里云"]="223.5.5.5,223.6.6.6"
DNS_PRESETS["114DNS"]="114.114.114.114,114.114.115.115"

# 测试DNS连通性
test_dns_connectivity() {
    local dns_server=$1
    echo -e "${BLUE}正在测试DNS服务器 $dns_server 的连通性...${NC}"
    
    # 使用nslookup测试DNS解析
    if command -v nslookup >/dev/null 2>&1; then
        if timeout 5 nslookup baidu.com $dns_server >/dev/null 2>&1; then
            echo -e "${GREEN}DNS服务器 $dns_server 连通性测试通过${NC}"
            return 0
        else
            echo -e "${RED}DNS服务器 $dns_server 连通性测试失败${NC}"
            return 1
        fi
    else
        # 如果没有nslookup，使用dig
        if command -v dig >/dev/null 2>&1; then
            if timeout 5 dig @$dns_server baidu.com >/dev/null 2>&1; then
                echo -e "${GREEN}DNS服务器 $dns_server 连通性测试通过${NC}"
                return 0
            else
                echo -e "${RED}DNS服务器 $dns_server 连通性测试失败${NC}"
                return 1
            fi
        else
            # 如果没有nslookup和dig，使用ping测试
            if command -v ping >/dev/null 2>&1; then
                echo -e "${YELLOW}使用ping测试DNS服务器连通性...${NC}"
                if ping -c 3 -W 3 $dns_server >/dev/null 2>&1; then
                    echo -e "${GREEN}DNS服务器 $dns_server ping测试通过${NC}"
                    return 0
                else
                    echo -e "${RED}DNS服务器 $dns_server ping测试失败${NC}"
                    return 1
                fi
            else
                echo -e "${YELLOW}警告: 系统缺少DNS测试工具(nslookup/dig/ping)，跳过连通性测试${NC}"
                return 0
            fi
        fi
    fi
}

# 应用DNS配置
apply_dns_config() {
    local primary_dns=$1
    local secondary_dns=$2
    
    echo -e "${BLUE}正在应用DNS配置...${NC}"
    
    # 备份原有配置
    if [ -f /etc/resolv.conf ]; then
        cp /etc/resolv.conf /etc/resolv.conf.backup.$(date +%Y%m%d_%H%M%S)
        echo -e "${GREEN}已备份原有DNS配置${NC}"
    fi
    
    # 写入新的DNS配置
    {
        echo "# DNS配置 - 由脚本自动生成于 $(date)"
        echo "nameserver $primary_dns"
        if [ -n "$secondary_dns" ]; then
            echo "nameserver $secondary_dns"
        fi
    } > /etc/resolv.conf
    
    echo -e "${GREEN}DNS配置已成功应用${NC}"
    echo -e "主DNS: $primary_dns"
    if [ -n "$secondary_dns" ]; then
        echo -e "备用DNS: $secondary_dns"
    fi
}

# 获取用户自定义DNS输入
get_custom_dns() {
    local primary_dns
    local secondary_dns
    
    while true; do
        echo -e "${YELLOW}请输入主DNS服务器地址:${NC}"
        read -p "主DNS: " primary_dns
        
        # 验证IP地址格式
        if [[ $primary_dns =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ ]]; then
            # 测试连通性
            if test_dns_connectivity "$primary_dns"; then
                break
            else
                echo -e "${RED}主DNS服务器连通性测试失败，请重新输入${NC}"
                continue
            fi
        else
            echo -e "${RED}无效的IP地址格式，请重新输入${NC}"
            continue
        fi
    done
    
    echo -e "${YELLOW}请输入备用DNS服务器地址(可选，直接回车跳过):${NC}"
    read -p "备用DNS: " secondary_dns
    
    if [ -n "$secondary_dns" ]; then
        if [[ $secondary_dns =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ ]]; then
            if ! test_dns_connectivity "$secondary_dns"; then
                echo -e "${RED}备用DNS服务器连通性测试失败，将只使用主DNS${NC}"
                secondary_dns=""
            fi
        else
            echo -e "${RED}无效的备用DNS地址格式，将只使用主DNS${NC}"
            secondary_dns=""
        fi
    fi
    
    apply_dns_config "$primary_dns" "$secondary_dns"
}

# 主DNS配置函数
configure_dns() {
    echo -e "${BLUE}=== DNS配置工具 ===${NC}"
    echo "请选择DNS配置选项:"
    
    local options=("谷歌" "Cloudflare" "腾讯云" "阿里云" "114DNS" "自定义DNS")
    local i=1
    
    for option in "${options[@]}"; do
        echo "$i) $option"
        ((i++))
    done
    
    echo "0) 返回主菜单"
    echo
    
    while true; do
        read -p "请输入选项编号: " choice
        
        case $choice in
            0)
                return
                ;;
            1|2|3|4|5)
                local selected_option="${options[$((choice-1))]}"
                local dns_config="${DNS_PRESETS[$selected_option]}"
                local primary_dns=$(echo $dns_config | cut -d',' -f1)
                local secondary_dns=$(echo $dns_config | cut -d',' -f2)
                
                echo -e "${BLUE}您选择了: $selected_option${NC}"
                echo "主DNS: $primary_dns"
                echo "备用DNS: $secondary_dns"
                
                # 测试连通性
                if test_dns_connectivity "$primary_dns" && test_dns_connectivity "$secondary_dns"; then
                    echo -e "${GREEN}DNS连通性测试通过，正在应用配置...${NC}"
                    apply_dns_config "$primary_dns" "$secondary_dns"
                    echo -e "${GREEN}DNS配置完成！${NC}"
                    read -p "按回车键继续..."
                    return
                else
                    echo -e "${RED}DNS连通性测试失败，无法应用此配置${NC}"
                    read -p "按回车键继续..."
                    continue
                fi
                ;;
            6)
                echo -e "${BLUE}您选择了: 自定义DNS${NC}"
                get_custom_dns
                echo -e "${GREEN}自定义DNS配置完成！${NC}"
                read -p "按回车键继续..."
                return
                ;;
            *)
                echo -e "${RED}无效选项，请重新输入${NC}"
                ;;
        esac
    done
}

# 如果脚本被直接执行，则运行主函数
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    # 检查是否为root用户
    if [ "$EUID" -ne 0 ]; then
        echo -e "${RED}此脚本需要root权限运行${NC}"
        echo "请使用: sudo $0"
        exit 1
    fi
    
    configure_dns
fi